'''
https://leetcode.cn/problems/number-of-ways-to-wear-different-hats-to-each-other/description/
'''
from collections import defaultdict
from functools import cache
from typing import List


class Solution:

    # 正向：以帽子做状态，根据人来选帽子
    def numberWays(self, hats: List[List[int]]) -> int:
        MOD = 10 ** 9 + 7
        n = len(hats)

        @cache
        def f(status, i):
            if i == n: return 1
            res = 0
            for hat in hats[i]:
                if status >> hat & 1: continue
                res += f(status | (1 << hat), i + 1) % MOD
            return res % MOD

        return f(0, 0)

    # 反向：以人来做状态，根据帽子来选人
    #   构建一个反向表，帽子可以选择那些人
    #   离散化，还要讲帽子离散化到0~39, 不然递归
    def numberWays2(self, hats: List[List[int]]) -> int:
        MOD = 10 ** 9 + 7
        hats_id = {}
        hats_people = defaultdict(list)
        id = 0
        for people, favorite_hats in enumerate(hats):
            for hat in favorite_hats:
                if hat not in hats_id:
                    hats_id[hat] = id
                    id += 1
                hats_people[hats_id[hat]].append(people)
        m, n = len(hats), id  # m: 多少人[0, m-1]， n多少帽子:[0, n-1]
        final_status = (1 << m) - 1  # 某个人是否带有帽子: status[i] == 0没带 else 带了

        @cache
        def f(status, i):
            if status == final_status: return 1
            if i == n: return 0
            res = f(status, i + 1) % MOD  # 不要当前帽子
            for people in hats_people[i]:
                if status >> people & 1: continue
                res += f(status | (1 << people), i + 1) % MOD  # 要当前帽子
            return res % MOD

        return f(0, 0)

    # 反向打表
    def numberWays3(self, hats: List[List[int]]) -> int:
        MOD = 10 ** 9 + 7
        hats_id = {}
        hats_people = defaultdict(list)
        id = 0
        for people, favorite_hats in enumerate(hats):
            for hat in favorite_hats:
                if hat not in hats_id:
                    hats_id[hat] = id
                    id += 1
                hats_people[hats_id[hat]].append(people)
        m, n = len(hats), id  # m: 多少人[0, m-1]， n多少帽子:[0, n-1]
        final_status = (1 << m) - 1  # 某个人是否带有帽子: status[i] == 0没带 else 带了
        dp = [[0] * (n + 1) for _ in range(final_status + 1)]
        # 第一维度依赖后边的，第二维度依赖后边的
        for i in range(n + 1):
            dp[final_status][i] = 1
        for status in range(final_status - 1, -1, -1):
            for i in range(n - 1, -1, -1):
                res = dp[status][i + 1] % MOD  # 不要当前帽子
                for people in hats_people[i]:
                    if status >> people & 1: continue
                    res += dp[status | (1 << people)][i + 1] % MOD  # 要当前帽子
                dp[status][i] = res % MOD
        return dp[0][0]


hats = [[3, 5, 1], [3, 5]]
print(Solution().numberWays2(hats))
